'''
Author: xialang333@gmail.com
Date: 2025-08-25 21:33:06
LastEditTime: 2025-08-26 09:08:05
Description: 
'''
import os
import sys
import json
from typing import Dict, List, Any
from datetime import datetime
from pydantic import BaseModel, Field
from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
from fastapi.security import HTTPBearer
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi import Request, Form
from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, Text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
import logging
import aiohttp

# 获取正确的资源路径
def get_resource_path(relative_path):
    """获取资源文件的绝对路径"""
    try:
        # 打包后的路径
        base_path = sys._MEIPASS
    except AttributeError:
        # 开发环境的路径
        base_path = os.path.dirname(os.path.abspath(__file__))
    
    path = os.path.join(base_path, relative_path)
    
    # 如果路径不存在，尝试在当前目录查找
    if not os.path.exists(path):
        path = os.path.join(os.getcwd(), relative_path)
    
    return path

# 数据库配置
SQLALCHEMY_DATABASE_URL = "sqlite:///./app.db"
engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

# 数据库模型
class Collector(Base):
    __tablename__ = "collectors"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, unique=True, index=True)
    type = Column(String)  # node-exporter, cadvisor, etc.
    url = Column(String)   # 采集器地址
    enabled = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class AlertRule(Base):
    __tablename__ = "alert_rules"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, unique=True, index=True)
    expr = Column(String)  # PromQL 表达式
    duration = Column(String)  # 持续时间
    severity = Column(String)  # 严重程度
    summary = Column(String)
    description = Column(String)
    enabled = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class NotificationConfig(Base):
    __tablename__ = "notification_configs"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, unique=True, index=True)
    type = Column(String)  # email, dingtalk, feishu, wechat
    config = Column(Text)  # JSON 配置
    enabled = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class AlertHistory(Base):
    __tablename__ = "alert_history"
    
    id = Column(Integer, primary_key=True, index=True)
    alert_name = Column(String)
    severity = Column(String)
    status = Column(String)  # firing, resolved
    summary = Column(String)
    description = Column(String)
    instance = Column(String)
    value = Column(String)
    notified = Column(Boolean, default=False)
    created_at = Column(DateTime, default=datetime.utcnow)

# 创建数据库表
Base.metadata.create_all(bind=engine)

# Pydantic 模型
class CollectorCreate(BaseModel):
    name: str
    type: str
    url: str
    enabled: bool = True

class CollectorResponse(BaseModel):
    id: int
    name: str
    type: str
    url: str
    enabled: bool
    created_at: datetime
    updated_at: datetime

    class Config:
        orm_mode = True

class AlertRuleCreate(BaseModel):
    name: str
    expr: str
    duration: str
    severity: str
    summary: str
    description: str
    enabled: bool = True

class AlertRuleResponse(BaseModel):
    id: int
    name: str
    expr: str
    duration: str
    severity: str
    summary: str
    description: str
    enabled: bool
    created_at: datetime
    updated_at: datetime

    class Config:
        orm_mode = True

class NotificationConfigCreate(BaseModel):
    name: str
    type: str
    config: Dict[str, Any]
    enabled: bool = True

class NotificationConfigResponse(BaseModel):
    id: int
    name: str
    type: str
    config: Dict[str, Any]
    enabled: bool
    created_at: datetime
    updated_at: datetime

    class Config:
        orm_mode = True

class AlertHistoryResponse(BaseModel):
    id: int
    alert_name: str
    severity: str
    status: str
    summary: str
    description: str
    instance: str
    value: str
    notified: bool
    created_at: datetime

    class Config:
        orm_mode = True


# 应用初始化
app = FastAPI(title="Exporter Manager API", description="管理 Prometheus 采集器和告警通知", version="1.0.0")
# 设置静态文件和模板路径
static_path = get_resource_path("static")
templates_path = get_resource_path("templates")
# 挂载静态文件和模板
# app.mount("/static", StaticFiles(directory="static"), name="static")
# templates = Jinja2Templates(directory="templates")
app.mount("/static", StaticFiles(directory=static_path), name="static")
templates = Jinja2Templates(directory=templates_path)
# 安全配置
security = HTTPBearer()

# 获取数据库会话
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# 工具函数
async def send_email_notification(config: Dict, message: str, subject: str):
    """发送邮件通知"""
    try:
        msg = MIMEMultipart()
        msg['From'] = config['smtp_user']
        msg['To'] = config['to']
        msg['Subject'] = subject
        
        msg.attach(MIMEText(message, 'plain'))
        
        server = smtplib.SMTP(config['smtp_server'], config['smtp_port'])
        server.starttls()
        server.login(config['smtp_user'], config['smtp_password'])
        server.send_message(msg)
        server.quit()
        return True
    except Exception as e:
        logging.error(f"发送邮件失败: {e}")
        return False

async def send_dingtalk_notification(config: Dict, message: str):
    """发送钉钉通知"""
    try:
        webhook_url = config['webhook_url']
        data = {
            "msgtype": "text",
            "text": {
                "content": message
            }
        }
        async with aiohttp.ClientSession() as session:
            async with session.post(webhook_url, json=data) as response:
                return response.status == 200
    except Exception as e:
        logging.error(f"发送钉钉通知失败: {e}")
        return False

async def send_feishu_notification(config: Dict, message: str):
    """发送飞书通知"""
    try:
        webhook_url = config['webhook_url']
        data = {
            "msg_type": "text",
            "content": {
                "text": message
            }
        }
        async with aiohttp.ClientSession() as session:
            async with session.post(webhook_url, json=data) as response:
                return response.status == 200
    except Exception as e:
        logging.error(f"发送飞书通知失败: {e}")
        return False

async def send_wechat_notification(config: Dict, message: str):
    """发送企业微信通知"""
    try:
        webhook_url = config['webhook_url']
        data = {
            "msgtype": "text",
            "text": {
                "content": message
            }
        }
        async with aiohttp.ClientSession() as session:
            async with session.post(webhook_url, json=data) as response:
                return response.status == 200
    except Exception as e:
        logging.error(f"发送企业微信通知失败: {e}")
        return False

async def send_notification(notification_type: str, config: Dict, message: str, subject: str = "告警通知"):
    """发送通知"""
    if notification_type == "email":
        return await send_email_notification(config, message, subject)
    elif notification_type == "dingtalk":
        return await send_dingtalk_notification(config, message)
    elif notification_type == "feishu":
        return await send_feishu_notification(config, message)
    elif notification_type == "wechat":
        return await send_wechat_notification(config, message)
    else:
        logging.error(f"不支持的通知类型: {notification_type}")
        return False

# 路由定义
@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request):
    """Web 界面首页"""
    return templates.TemplateResponse("index.html", {"request": request})

# 采集器管理 API
@app.post("/api/collectors/", response_model=CollectorResponse)
async def create_collector(collector: CollectorCreate, db: Session = Depends(get_db)):
    """创建采集器"""
    db_collector = Collector(**collector.dict())
    db.add(db_collector)
    db.commit()
    db.refresh(db_collector)
    return db_collector

@app.get("/api/collectors/", response_model=List[CollectorResponse])
async def list_collectors(db: Session = Depends(get_db)):
    """获取采集器列表"""
    return db.query(Collector).all()

@app.get("/api/collectors/{collector_id}", response_model=CollectorResponse)
async def get_collector(collector_id: int, db: Session = Depends(get_db)):
    """获取特定采集器"""
    collector = db.query(Collector).filter(Collector.id == collector_id).first()
    if not collector:
        raise HTTPException(status_code=404, detail="采集器不存在")
    return collector

@app.put("/api/collectors/{collector_id}", response_model=CollectorResponse)
async def update_collector(collector_id: int, collector: CollectorCreate, db: Session = Depends(get_db)):
    """更新采集器"""
    db_collector = db.query(Collector).filter(Collector.id == collector_id).first()
    if not db_collector:
        raise HTTPException(status_code=404, detail="采集器不存在")
    
    for key, value in collector.dict().items():
        setattr(db_collector, key, value)
    
    db.commit()
    db.refresh(db_collector)
    return db_collector

@app.delete("/api/collectors/{collector_id}")
async def delete_collector(collector_id: int, db: Session = Depends(get_db)):
    """删除采集器"""
    collector = db.query(Collector).filter(Collector.id == collector_id).first()
    if not collector:
        raise HTTPException(status_code=404, detail="采集器不存在")
    
    db.delete(collector)
    db.commit()
    return {"message": "采集器已删除"}

# 告警规则管理 API
@app.post("/api/alert_rules/", response_model=AlertRuleResponse)
async def create_alert_rule(alert_rule: AlertRuleCreate, db: Session = Depends(get_db)):
    """创建告警规则"""
    db_alert_rule = AlertRule(**alert_rule.dict())
    db.add(db_alert_rule)
    db.commit()
    db.refresh(db_alert_rule)
    return db_alert_rule

@app.get("/api/alert_rules/", response_model=List[AlertRuleResponse])
async def list_alert_rules(db: Session = Depends(get_db)):
    """获取告警规则列表"""
    return db.query(AlertRule).all()

@app.get("/api/alert_rules/{alert_rule_id}", response_model=AlertRuleResponse)
async def get_alert_rule(alert_rule_id: int, db: Session = Depends(get_db)):
    """获取特定告警规则"""
    alert_rule = db.query(AlertRule).filter(AlertRule.id == alert_rule_id).first()
    if not alert_rule:
        raise HTTPException(status_code=404, detail="告警规则不存在")
    return alert_rule

@app.put("/api/alert_rules/{alert_rule_id}", response_model=AlertRuleResponse)
async def update_alert_rule(alert_rule_id: int, alert_rule: AlertRuleCreate, db: Session = Depends(get_db)):
    """更新告警规则"""
    db_alert_rule = db.query(AlertRule).filter(AlertRule.id == alert_rule_id).first()
    if not db_alert_rule:
        raise HTTPException(status_code=404, detail="告警规则不存在")
    
    for key, value in alert_rule.dict().items():
        setattr(db_alert_rule, key, value)
    
    db.commit()
    db.refresh(db_alert_rule)
    return db_alert_rule

@app.delete("/api/alert_rules/{alert_rule_id}")
async def delete_alert_rule(alert_rule_id: int, db: Session = Depends(get_db)):
    """删除告警规则"""
    alert_rule = db.query(AlertRule).filter(AlertRule.id == alert_rule_id).first()
    if not alert_rule:
        raise HTTPException(status_code=404, detail="告警规则不存在")
    
    db.delete(alert_rule)
    db.commit()
    return {"message": "告警规则已删除"}

# 通知配置管理 API
@app.post("/api/notification_configs/", response_model=NotificationConfigResponse)
async def create_notification_config(notification_config: NotificationConfigCreate, db: Session = Depends(get_db)):
    """创建通知配置"""
    db_notification_config = NotificationConfig(
        name=notification_config.name,
        type=notification_config.type,
        config=json.dumps(notification_config.config),
        enabled=notification_config.enabled
    )
    db.add(db_notification_config)
    db.commit()
    db.refresh(db_notification_config)
    return db_notification_config

@app.get("/api/notification_configs/", response_model=List[NotificationConfigResponse])
async def list_notification_configs(db: Session = Depends(get_db)):
    """获取通知配置列表"""
    configs = db.query(NotificationConfig).all()
    for config in configs:
        config.config = json.loads(config.config)
    return configs

@app.get("/api/notification_configs/{config_id}", response_model=NotificationConfigResponse)
async def get_notification_config(config_id: int, db: Session = Depends(get_db)):
    """获取特定通知配置"""
    notification_config = db.query(NotificationConfig).filter(NotificationConfig.id == config_id).first()
    if not notification_config:
        raise HTTPException(status_code=404, detail="通知配置不存在")
    
    notification_config.config = json.loads(notification_config.config)
    return notification_config

@app.put("/api/notification_configs/{config_id}", response_model=NotificationConfigResponse)
async def update_notification_config(config_id: int, notification_config: NotificationConfigCreate, db: Session = Depends(get_db)):
    """更新通知配置"""
    db_notification_config = db.query(NotificationConfig).filter(NotificationConfig.id == config_id).first()
    if not db_notification_config:
        raise HTTPException(status_code=404, detail="通知配置不存在")
    
    db_notification_config.name = notification_config.name
    db_notification_config.type = notification_config.type
    db_notification_config.config = json.dumps(notification_config.config)
    db_notification_config.enabled = notification_config.enabled
    
    db.commit()
    db.refresh(db_notification_config)
    db_notification_config.config = json.loads(db_notification_config.config)
    return db_notification_config

@app.delete("/api/notification_configs/{config_id}")
async def delete_notification_config(config_id: int, db: Session = Depends(get_db)):
    """删除通知配置"""
    notification_config = db.query(NotificationConfig).filter(NotificationConfig.id == config_id).first()
    if not notification_config:
        raise HTTPException(status_code=404, detail="通知配置不存在")
    
    db.delete(notification_config)
    db.commit()
    return {"message": "通知配置已删除"}

# Prometheus 指标端点
@app.get("/metrics")
async def metrics(db: Session = Depends(get_db)):
    """提供 Prometheus 可采集的指标"""
    # 这里应该实现从各个采集器收集指标的逻辑
    # 简化版本：返回系统状态指标
    collectors_count = db.query(Collector).count()
    enabled_collectors_count = db.query(Collector).filter(Collector.enabled == True).count()
    alert_rules_count = db.query(AlertRule).count()
    enabled_alert_rules_count = db.query(AlertRule).filter(AlertRule.enabled == True).count()
    
    metrics_data = f"""
# HELP prometheus_manager_collectors_total Total number of collectors
# TYPE prometheus_manager_collectors_total gauge
prometheus_manager_collectors_total {collectors_count}

# HELP prometheus_manager_enabled_collectors_total Total number of enabled collectors
# TYPE prometheus_manager_enabled_collectors_total gauge
prometheus_manager_enabled_collectors_total {enabled_collectors_count}

# HELP prometheus_manager_alert_rules_total Total number of alert rules
# TYPE prometheus_manager_alert_rules_total gauge
prometheus_manager_alert_rules_total {alert_rules_count}

# HELP prometheus_manager_enabled_alert_rules_total Total number of enabled alert rules
# TYPE prometheus_manager_enabled_alert_rules_total gauge
prometheus_manager_enabled_alert_rules_total {enabled_alert_rules_count}
"""
    return metrics_data

# 告警接收端点
@app.post("/alerts")
async def receive_alerts(request: Request, background_tasks: BackgroundTasks, db: Session = Depends(get_db)):
    """接收 Prometheus 发送的告警"""
    try:
        alert_data = await request.json()
        logging.info(f"收到告警: {alert_data}")
        
        for alert in alert_data.get('alerts', []):
            # 保存告警历史
            alert_history = AlertHistory(
                alert_name=alert['labels'].get('alertname', 'unknown'),
                severity=alert['labels'].get('severity', 'unknown'),
                status=alert['status'],
                summary=alert['annotations'].get('summary', ''),
                description=alert['annotations'].get('description', ''),
                instance=alert['labels'].get('instance', ''),
                value=alert['labels'].get('value', ''),
                notified=False
            )
            db.add(alert_history)
            db.commit()
            
            # 如果是 firing 状态，发送通知
            if alert['status'] == 'firing':
                background_tasks.add_task(send_alert_notifications, alert, db)
        
        return {"status": "success"}
    except Exception as e:
        logging.error(f"处理告警时出错: {e}")
        raise HTTPException(status_code=500, detail="处理告警时出错")

async def send_alert_notifications(alert: Dict, db: Session):
    """发送告警通知"""
    try:
        # 获取所有启用的通知配置
        notification_configs = db.query(NotificationConfig).filter(NotificationConfig.enabled == True).all()
        
        # 构建消息内容
        alert_name = alert['labels'].get('alertname', '未知告警')
        severity = alert['labels'].get('severity', '未知')
        summary = alert['annotations'].get('summary', '')
        description = alert['annotations'].get('description', '')
        instance = alert['labels'].get('instance', '')
        
        message = f"""
告警名称: {alert_name}
严重程度: {severity}
实例: {instance}
摘要: {summary}
描述: {description}
状态: {alert['status']}
开始时间: {alert['startsAt']}
"""
        subject = f"[{severity}] {alert_name} 告警"
        
        # 发送通知
        for config in notification_configs:
            config_dict = json.loads(config.config)
            success = await send_notification(config.type, config_dict, message, subject)
            
            if success:
                logging.info(f"通过 {config.type} 发送告警通知成功")
            else:
                logging.error(f"通过 {config.type} 发送告警通知失败")
        
        # 更新告警历史记录为已通知
        # 这里简化处理，实际应该根据具体告警更新对应的记录
        return True
    except Exception as e:
        logging.error(f"发送告警通知时出错: {e}")
        return False

# 获取告警历史
@app.get("/api/alert_history/", response_model=List[AlertHistoryResponse])
async def get_alert_history(limit: int = 100, db: Session = Depends(get_db)):
    """获取告警历史"""
    return db.query(AlertHistory).order_by(AlertHistory.created_at.desc()).limit(limit).all()

# Web 界面路由
@app.get("/collectors", response_class=HTMLResponse)
async def collectors_ui(request: Request, db: Session = Depends(get_db)):
    """采集器管理界面"""
    collectors = db.query(Collector).all()
    return templates.TemplateResponse("collectors.html", {"request": request, "collectors": collectors})

@app.get("/alert_rules", response_class=HTMLResponse)
async def alert_rules_ui(request: Request, db: Session = Depends(get_db)):
    """告警规则管理界面"""
    alert_rules = db.query(AlertRule).all()
    return templates.TemplateResponse("alert_rules.html", {"request": request, "alert_rules": alert_rules})

@app.get("/notifications", response_class=HTMLResponse)
async def notifications_ui(request: Request, db: Session = Depends(get_db)):
    """通知配置管理界面"""
    notification_configs = db.query(NotificationConfig).all()
    for config in notification_configs:
        config.config = json.loads(config.config)
    return templates.TemplateResponse("notifications.html", {"request": request, "notification_configs": notification_configs})

@app.get("/alert_history", response_class=HTMLResponse)
async def alert_history_ui(request: Request, db: Session = Depends(get_db)):
    """告警历史界面"""
    alert_history = db.query(AlertHistory).order_by(AlertHistory.created_at.desc()).limit(100).all()
    return templates.TemplateResponse("alert_history.html", {"request": request, "alert_history": alert_history})

# 启动应用
if __name__ == "__main__":
    import uvicorn
    import argparse
    parser = argparse.ArgumentParser(description="FastAPI Server")
    parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind")
    parser.add_argument("--port", type=int, default=8000, help="Port to listen on")
    parser.add_argument("--reload", action="store_true", help="Enable reload")
    parser.add_argument("--workers", type=int, default=1,help="Number of worker processes (default: 1)")
    
    args = parser.parse_args()
    print("=" * 50)
    print(f"🚀 Starting FastAPI Server")
    print(f"📍 Host: {args.host}")
    print(f"🔌 Port: {args.port}")
    print(f"🔄 Reload: {args.reload}")
    print(f"👥 Workers: {args.workers}")
    print("=" * 50)
    uvicorn.run(app,
        host=args.host,
        port=args.port,
        reload=args.reload,
        workers=args.workers,
        log_level="info")